﻿// 2867. 统计树中的合法路径数目.cpp : 此文件包含 "main" 函数。程序执行将在此处开始并结束。
//

#include <iostream>
#include <vector>
#include <unordered_set>


using namespace std;

/*
https://leetcode.cn/problems/count-valid-paths-in-a-tree/description/

给你一棵 n 个节点的无向树，节点编号为 1 到 n 。
给你一个整数 n 和一个长度为 n - 1 的二维整数数组 edges ，其中 edges[i] = [ui, vi] 表示节点 ui 和 vi 在树中有一条边。

请你返回树中的 合法路径数目 。

如果在节点 a 到节点 b 之间 恰好有一个 节点的编号是质数，那么我们称路径 (a, b) 是 合法的 。

注意：

路径 (a, b) 指的是一条从节点 a 开始到节点 b 结束的一个节点序列，序列中的节点 互不相同 ，且相邻节点之间在树上有一条边。
路径 (a, b) 和路径 (b, a) 视为 同一条 路径，且只计入答案 一次 。


示例 1：
输入：n = 5, edges = [[1,2],[1,3],[2,4],[2,5]]
输出：4
解释：恰好有一个质数编号的节点路径有：
- (1, 2) 因为路径 1 到 2 只包含一个质数 2 。
- (1, 3) 因为路径 1 到 3 只包含一个质数 3 。
- (1, 4) 因为路径 1 到 4 只包含一个质数 2 。
- (2, 4) 因为路径 2 到 4 只包含一个质数 2 。
只有 4 条合法路径。


示例 2：
输入：n = 6, edges = [[1,2],[1,3],[2,4],[3,5],[3,6]]
输出：6
解释：恰好有一个质数编号的节点路径有：
- (1, 2) 因为路径 1 到 2 只包含一个质数 2 。
- (1, 3) 因为路径 1 到 3 只包含一个质数 3 。
- (1, 4) 因为路径 1 到 4 只包含一个质数 2 。
- (1, 6) 因为路径 1 到 6 只包含一个质数 3 。
- (2, 4) 因为路径 2 到 4 只包含一个质数 2 。
- (3, 6) 因为路径 3 到 6 只包含一个质数 3 。
只有 6 条合法路径。


提示：
1 <= n <= 105
edges.length == n - 1
edges[i].length == 2
1 <= ui, vi <= n
输入保证 edges 形成一棵合法的树。
*/


class Solution {
public:
    int h[100010], ne[200010], e[200010], idx = 0;
    long long dp[100010][2];
    unordered_set<int> ss;
    long long ans = 0;
    void get_primes(int n)
    {
        int cnt = 0;     // primes[]存储所有素数
        vector<int> primes(n+10);
        vector<bool> st(n+10);         // st[x]存储x是否被筛掉
        for (int i = 2; i <= n; i++)
        {
            if (!st[i]) { primes[cnt++] = i;  ss.insert(i); }
            for (int j = 0; primes[j] <= n / i; j++)
            {
                st[primes[j] * i] = true;
                if (i % primes[j] == 0) break;
            }
        }
    }
    void add(int a, int b) {
        e[idx] = b, ne[idx] = h[a], h[a] = idx++;
    }

    void dfs(int x,int father) {
        vector<int> next;
        for (int i = h[x]; i != -1; i = ne[i]) {
            int j = e[i];
            if (father == j) continue;
            dfs(j,x);
            next.push_back(j);
        }

        if (ss.count(x) != 0) {
            dp[x][1] = 1;
        }
        else {
            dp[x][0] = 1;
        }

        if (!next.empty()) {
            //dp表示是以该点为起点的  合法路径和无质数路径数目
            long long validPath = 0;
            long long noPrimPath = 0;
            for (int i = 0; i < next.size(); i++) {
                int p = next[i];
                validPath += dp[p][1];
                noPrimPath += dp[p][0];
            }
            if (ss.count(x) == 1) {                                                                                                                                                                                                                                                                         
                dp[x][1] += noPrimPath;
                dp[x][0] = 0;   //起点是质数 不可能存在无质数路径
                //再来看 经过点x的合法路径  x本身是质数 那么任意两个子节点的无质数路径乘积就是一个合法路径
                long long curr = 0;
                for (int i = 0; i < next.size(); i++) {
                    int p = next[i];
                    curr += dp[p][0] * (noPrimPath - dp[p][0]);
                }

                //由于 ab  ba 经过当前点的路径重复了2次
                curr /= 2;
                //cout << " middle = " << x << ". path = " << curr << endl;
                ans += curr;
            }
            else {
                dp[x][1] += validPath;
                dp[x][0] += noPrimPath;
                //再来看 经过点x的合法路径 x不是质数 那么任意两个子节点的无质数路径和合法路径乘积就是一个合法路径
                long long curr = 0;
                for (int i = 0; i < next.size(); i++) {
                    int p = next[i];
                    curr += dp[p][1] * (noPrimPath - dp[p][0]);
                }
                //cout << " middle = " << x << ". path = " << curr << endl;
                ans += curr;
            }
        }

        return ;
    }
    long long countPaths(int n, vector<vector<int>>& edges) {
        get_primes(n+10);
        memset(h, -1, sizeof h);
        memset(dp, 0, sizeof dp);

        for (int i = 0; i < edges.size(); i++) {
            int a = edges[i][0]; int b = edges[i][1];
            add(a, b); add(b, a);
        }
        dfs(1,-1);

        for (int i = 1; i <= n; i++) {
            if (ss.count(i)==0) {
                ans += dp[i][1];
            }
            else {
                ans += dp[i][1] - 1;
            }
        }

        return ans;
    }
};


int main()
{
    Solution s;
    vector<vector<int>> v{
        {1,2},{1,3},{2,4},{3,5},{3,6}
    };
    s.countPaths(6,v);
}

